import paddle
class CrossEntropyLoss(paddle.nn.Layer):
    def __init__(self,ignore_lable=255): #默认忽略255下标值不进行损失计算
        super(CrossEntropyLoss,self).__init__()
        self.loss = paddle.nn.CrossEntropyLoss(ignore_index=ignore_lable)
    def forward(self,pred,label,weight=None):
        assert pred.dim() == 4
        assert label.dim() == 3
        assert pred.shape[0] == label.shape[0], "{0} vs {1} ".format(pred.shape[0], label.shape[0])
        assert pred.shape[2] == label.shape[1], "{0} vs {1} ".format(pred.shape[2], label.shape[1])
        assert pred.shape[3] == label.shape[2], "{0} vs {1} ".format(pred.shape[3], label.shape[2])
        c = pred.shape[1]
        # [batch,c,h,w] -> [batch,h,w,c] -> [batch*h*w,c]
        pred = pred.transpose([0,2,3,1]).reshape([-1,c])
        # [batch,h,w] -> [batch*h*w,1]
        label = label.reshape([-1,1])
        return self.loss(pred,label) #忽略下标为255的loss计算